package com.jd.dlc.lock;

/**
 * @program: dlc
 * @ClassName KeyLock
 * @description:
 * @author: sizegang
 * @create: 2022-05-02
 **/

import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Semaphore;

public class KeyLock<K> {

    /**
     * Save all locked KEYs and their semaphores
     */
    private final ConcurrentMap<K, Semaphore> map = new ConcurrentHashMap<K, Semaphore>();
    /**
     * Save the KEY locked by each thread and its lock count
     */
    private final ThreadLocal<Map<K, LockInfo>> local = new ThreadLocal<Map<K, LockInfo>>() {
        @Override
        protected Map<K, LockInfo> initialValue() {
            return new HashMap<K, LockInfo>();
        }
    };

    /**
     * Lock the key, other threads waiting for this key will enter the wait until {@link #unlock(K)} is called
     * Use hashcode and equals to determine whether the keys are the same, so the keys must implement {@link #hashCode()} and
     * {@link #equals(Object)} method
     *
     * @param key
     */
    public void lock(K key) {
        if (key == null){
            return;
        }

        LockInfo info = local.get().get(key);
        if (info == null) {
            Semaphore current = new Semaphore(1);
            current.acquireUninterruptibly();
            Semaphore previous = map.put(key, current);
            if (previous != null){
                previous.acquireUninterruptibly();
            }

            local.get().put(key, new LockInfo(current));
        } else {
            info.lockCount++;
        }
    }

    /**
     * Release the key and wake up other threads waiting for this key
     * @param key
     */
    public void unlock(K key) {
        if (key == null){
            return;
        }

        LockInfo info = local.get().get(key);
        if (info != null && --info.lockCount == 0) {
            info.current.release();
            map.remove(key, info.current);
            local.get().remove(key);
        }
    }

    /**
     * Lock multiple keys
     * It is recommended to sort the keys before calling this method and use the same locking order to prevent deadlocks
     * @param keys
     */
    public void lock(K[] keys) {
        if (keys == null){
            return;
        }
        for (K key : keys) {
            lock(key);
        }
    }

    /**
     * release multiple keys
     * @param keys
     */
    public void unlock(K[] keys) {
        if (keys == null){
            return;
        }
        for (K key : keys) {
            unlock(key);
        }
    }

    private static class LockInfo {
        private final Semaphore current;
        private int lockCount;

        private LockInfo(Semaphore current) {
            this.current = current;
            this.lockCount = 1;
        }
    }
}
